
#!/usr/bin/env python3
"""
训练脚本示例
python train.py --data_dir ./ct_images --labels_csv labels.csv --model_out svm.pkl
"""
import argparse
import os
import csv
import joblib
from utils import preprocess_image, extract_features
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report

def load_dataset(data_dir, csv_path):
    """根据CSV文件读取影像路径与标签"""
    X, y = [], []
    with open(csv_path, newline='') as f:
        reader = csv.DictReader(f)
        for row in reader:
            fname = row['filename']
            label = int(row['label'])  # 0:良性  1:恶性
            img_path = os.path.join(data_dir, fname)
            img = preprocess_image(img_path)
            feat = extract_features(img)
            X.append(feat)
            y.append(label)
    return np.array(X), np.array(y)

if __name__ == '__main__':
    import numpy as np
    parser = argparse.ArgumentParser(description='训练肺结节SVM分类器')
    parser.add_argument('--data_dir', required=True, help='影像文件夹')
    parser.add_argument('--labels_csv', required=True, help='CSV文件：filename,label')
    parser.add_argument('--model_out', default='svm.pkl', help='输出模型路径')
    args = parser.parse_args()

    print("正在加载数据集...")
    X, y = load_dataset(args.data_dir, args.labels_csv)
    print(f"样本数: {len(y)}, 特征维度: {X.shape[1]}")

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

    print("正在训练SVM...")
    clf = SVC(kernel='rbf', probability=True, random_state=42)
    clf.fit(X_train, y_train)

    y_pred = clf.predict(X_test)
    print("验证集准确率:", accuracy_score(y_test, y_pred))
    print(classification_report(y_test, y_pred, target_names=['良性', '恶性']))

    joblib.dump(clf, args.model_out)
    print(f"模型已保存至 {args.model_out}")
